package igloo.aws;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.services.dynamodb.AmazonDynamoDBClient;
import com.amazonaws.services.dynamodb.model.BatchGetItemRequest;
import com.amazonaws.services.dynamodb.model.BatchGetItemResult;
import com.amazonaws.services.dynamodb.model.BatchWriteItemRequest;
import com.amazonaws.services.dynamodb.model.BatchWriteItemResult;
import com.amazonaws.services.dynamodb.model.CreateTableRequest;
import com.amazonaws.services.dynamodb.model.CreateTableResult;
import com.amazonaws.services.dynamodb.model.DeleteItemRequest;
import com.amazonaws.services.dynamodb.model.DeleteItemResult;
import com.amazonaws.services.dynamodb.model.DescribeTableRequest;
import com.amazonaws.services.dynamodb.model.DescribeTableResult;
import com.amazonaws.services.dynamodb.model.GetItemRequest;
import com.amazonaws.services.dynamodb.model.GetItemResult;
import com.amazonaws.services.dynamodb.model.KeysAndAttributes;
import com.amazonaws.services.dynamodb.model.ListTablesRequest;
import com.amazonaws.services.dynamodb.model.ListTablesResult;
import com.amazonaws.services.dynamodb.model.PutItemRequest;
import com.amazonaws.services.dynamodb.model.PutItemResult;
import com.amazonaws.services.dynamodb.model.QueryRequest;
import com.amazonaws.services.dynamodb.model.QueryResult;
import com.amazonaws.services.dynamodb.model.UpdateItemRequest;
import com.amazonaws.services.dynamodb.model.UpdateItemResult;
import com.amazonaws.services.dynamodb.model.WriteRequest;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;

public class IglooDynamoClient
{
    private AWSCredentials mCredentials;

    private AmazonDynamoDBClient mClient;
    private ExecutorService mExecutor;

    public IglooDynamoClient(AWSCredentials creds, ExecutorService executor)
    {
        mCredentials = creds;
        mClient = new AmazonDynamoDBClient(mCredentials);

        mExecutor = executor;
    }

    public AmazonDynamoDBClient getClient()
    {
        return mClient;
    }

    // 5 writes/s free
    // 10 reads/s free
    public ListenableFuture<CreateTableResult> createTable(final CreateTableRequest request)
    {
        final SettableFuture<CreateTableResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    CreateTableResult result = mClient.createTable(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<Boolean> hasTable(final String tableName)
    {
        final SettableFuture<Boolean> future = SettableFuture.create();

        ListenableFuture<List<String>> tablesFuture = listAllTables();
        Futures.addCallback(tablesFuture, new FutureCallback<List<String>>() {
            public void onSuccess(List<String> tables)
            {
                future.set(tables.contains(tableName));
            }

            public void onFailure(Throwable t)
            {
                future.setException(t);
            }
        }, mExecutor);

        return future;
    }

    public ListenableFuture<Map<String, Boolean>> hasTables(final List<String> tableNames)
    {
        final SettableFuture<Map<String, Boolean>> future = SettableFuture.create();

        ListenableFuture<List<String>> tablesFuture = listAllTables();
        Futures.addCallback(tablesFuture, new FutureCallback<List<String>>() {
            public void onSuccess(List<String> tables)
            {
                Map<String, Boolean> map = new HashMap<String, Boolean>();
                for(String tableName : tableNames)
                    map.put(tableName, tables.contains(tableName));

                future.set(map);
            }

            public void onFailure(Throwable t)
            {
                future.setException(t);
            }
        }, mExecutor);

        return future;
    }

    public ListenableFuture<DescribeTableResult> describeTable(final DescribeTableRequest request)
    {
        final SettableFuture<DescribeTableResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    DescribeTableResult result = mClient.describeTable(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<GetItemResult> getItem(final GetItemRequest request)
    {
        final SettableFuture<GetItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    GetItemResult result = mClient.getItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<BatchGetItemResult> batchGetItem(final BatchGetItemRequest request)
    {
        final SettableFuture<BatchGetItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    BatchGetItemResult result = mClient.batchGetItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<PutItemResult> putItem(final PutItemRequest request)
    {
        final SettableFuture<PutItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    PutItemResult result = mClient.putItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<UpdateItemResult> updateItem(final UpdateItemRequest request)
    {
        final SettableFuture<UpdateItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    UpdateItemResult result = mClient.updateItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<BatchWriteItemResult> batchWriteItems(
            final BatchWriteItemRequest request)
    {
        final SettableFuture<BatchWriteItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    BatchWriteItemResult result = mClient.batchWriteItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<DeleteItemResult> deleteItem(final DeleteItemRequest request)
    {
        final SettableFuture<DeleteItemResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    DeleteItemResult result = mClient.deleteItem(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<QueryResult> query(final QueryRequest request)
    {
        final SettableFuture<QueryResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    QueryResult result = mClient.query(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<ListTablesResult> listTables()
    {
        final SettableFuture<ListTablesResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    ListTablesResult result = mClient.listTables();
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<ListTablesResult> listTables(final ListTablesRequest request)
    {
        final SettableFuture<ListTablesResult> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                try
                {
                    ListTablesResult result = mClient.listTables(request);
                    future.set(result);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }

    public ListenableFuture<List<String>> listAllTables()
    {
        final SettableFuture<List<String>> future = SettableFuture.create();

        mExecutor.execute(new Runnable() {
            @Override
            public void run()
            {
                List<String> tableNames = new LinkedList<String>();
                try
                {
                    ListTablesResult result = null;
                    ListTablesRequest request = null;
                    do
                    {
                        if(result != null)
                        {
                            request = new ListTablesRequest();
                            request.setExclusiveStartTableName(request.getExclusiveStartTableName());
                        }
                        else
                            request = new ListTablesRequest();

                        result = mClient.listTables(request);
                        tableNames.addAll(result.getTableNames());

                    } while(result != null && result.getLastEvaluatedTableName() != null);

                    future.set(tableNames);
                }
                catch(Exception e)
                {
                    future.setException(e);
                }
            }
        });

        return future;
    }
}
